package io.github.javpower.imagerex.service;

import lombok.extern.slf4j.Slf4j;
import org.datavec.image.loader.NativeImageLoader;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.transferlearning.TransferLearningHelper;
import org.deeplearning4j.zoo.ZooModel;
import org.deeplearning4j.zoo.model.ResNet50;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.VGG16ImagePreProcessor;
import org.springframework.stereotype.Service;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

@Slf4j
@Service
public class FeatureExtractor {

    private final ComputationGraph resNet50;
    private final TransferLearningHelper transferLearningHelper;
    private final NativeImageLoader loader;
    private final DataNormalization scaler;

    public FeatureExtractor() throws IOException {
        // 初始化图片加载器，只初始化一次
        loader = new NativeImageLoader(224, 224, 3);

        // 加载预训练ResNet50模型
        ZooModel zooModel = ResNet50.builder().build();
        resNet50 = (ComputationGraph) zooModel.initPretrained();

        // 查看模型层名确认输出层，这里一般取倒数第二层 "avg_pool"
        transferLearningHelper = new TransferLearningHelper(resNet50, "avg_pool");

        // ResNet50一般和VGG16一样用这个预处理器
        scaler = new VGG16ImagePreProcessor();

        log.info("ResNet50模型加载完成");
    }

    public INDArray extractFeatures(String imagePath) throws IOException {
        File file = new File(imagePath);
        try (InputStream inputStream = new FileInputStream(file)) {
            INDArray image = loader.asMatrix(inputStream);
            scaler.transform(image);
            DataSet dataSet = new DataSet(image, null);
            INDArray features = transferLearningHelper.featurize(dataSet).getFeatures();
            return features.reshape(features.size(1));
        }
    }

    public List<Float> getVectorFeatures(String path) {
        INDArray indArray;
        try {
            indArray = extractFeatures(path);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        float[] floatArray = indArray.toFloatVector();
        List<Float> floatList = new ArrayList<>();
        for (float f : floatArray) {
            floatList.add(f);
        }
        return floatList;
    }
}
